#include "meshFunctions.hpp"
#include "structMeshPart.hpp"
#include "structMesh_testFunctions.hpp"
#include "structMeshDistribution.hpp"

#define MESH_NAME "RectMesh"
#define MESH_TYPE RectMesh

#include "structMeshPartHelper.hpp"

std::array<Array2Xi, 6> sizesAndOffsetsForTrueRectMesh();

TEST_CASE( "Non-uniform RectMesh: Constructing StructMeshPart in different ways" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		Array2i globalMeshSize {10,10};

		bool locallyAssigned {false};

		auto test = [&]( const StructMeshPart<RectMesh>& part){
			if( !locallyAssigned ){
				REQUIRE( comm.size() == part.comm().size() );
				REQUIRE( comm.rank() == part.comm().rank() );

				REQUIRE( (part.meshSizeGlobal() == globalMeshSize).all() );
				Array2Xi patchLimits (2, 5);
				patchLimits <<
					0, 5, 110, 120, 100,
					4, 9, 209, 219, 109;
				REQUIRE( (part.boundaryPatchLimits() == patchLimits).all() );
			}
		
			if( comm.size() == 4) {
				REQUIRE( (part.meshSize() == Array2i{5,5}).all() );

				if( comm.rank() == 0 ) {
					REQUIRE( (part.offset()   == Array2i{0,0}).all() );
					REQUIRE(  part.neighbours(BoundaryLocation::top)[0] == OverlapSegment{ 2, 5, 0, 0 });
					REQUIRE( part.atBoundary(BoundaryLocation::bottom) );
					REQUIRE( part.atBoundary(BoundaryLocation::left));
				}
				if( comm.rank() == 1 ) {
					REQUIRE( (part.offset()   == Array2i{5,0}).all() );
					REQUIRE(  part.neighbours(BoundaryLocation::top)[0] == OverlapSegment{ 3, 5, 0, 0 });
					REQUIRE( part.atBoundary(BoundaryLocation::bottom) );
					REQUIRE( part.atBoundary(BoundaryLocation::right));
				}
				if( comm.rank() == 2 ) {
					REQUIRE( (part.offset()   == Array2i{0,5}).all() );
					REQUIRE(  part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{ 0, 5, 0, 0 });
					REQUIRE( part.atBoundary(BoundaryLocation::top) );
					REQUIRE( part.atBoundary(BoundaryLocation::left));
				}
				if( comm.rank() == 3 ) {
					REQUIRE( (part.offset()   == Array2i{5,5}).all() );
					REQUIRE(  part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{ 1, 5, 0, 0 });
					REQUIRE( part.atBoundary(BoundaryLocation::top) );
					REQUIRE( part.atBoundary(BoundaryLocation::right));
				}

				REQUIRE( (part.nParts()      == Array2i{2,2}).all() );
				REQUIRE(  part.band()        == comm.rank() / 2 );
				REQUIRE(  part.indexInBand() == comm.rank() % 2);
			}
		};

		auto addGlobalPatches = [&](RectMesh& mesh){
			addStandardBoundaryPatches(mesh);

			mesh.eraseBoundaryPatch("bottom");
			mesh.addBoundaryPatch("bottomLeft",  ArrayXi::LinSpaced(5, 0, 4));
			mesh.addBoundaryPatch("bottomRight", ArrayXi::LinSpaced(5, 5, 9));
		};

		auto addLocalPatches =[&](RectMesh& mesh){
			addStandardBoundaryPatches(mesh);

			if(comm.rank() == 0 || comm.rank() == 1)
				mesh.boundaryPatch("top")   = BoundaryPatch{"top"};
			if(comm.rank() == 1 || comm.rank() == 3)
				mesh.boundaryPatch("left")  = BoundaryPatch{"left"};
			if(comm.rank() == 0 || comm.rank() == 2)
				mesh.boundaryPatch("right") = BoundaryPatch{"right"};
			
			mesh.eraseBoundaryPatch("bottom");
			mesh.addBoundaryPatch("bottomLeft",  ArrayXi{} );
			mesh.addBoundaryPatch("bottomRight", ArrayXi{} );

			if( comm.rank() == 0 ){
				mesh.boundaryPatch("bottomLeft").setEdgesAndCells(mesh, ArrayXi::LinSpaced(3, 0, 2));
				mesh.boundaryPatch("bottomRight").setEdgesAndCells(mesh, ArrayXi::LinSpaced(2, 3, 4));
			}
			if( comm.rank() == 1){
				mesh.boundaryPatch("bottomRight").setEdgesAndCells(mesh, ArrayXi::LinSpaced(5,0,4));
			}
		};

		auto localOffset = [&]() -> Array2i{
			Array2i offset;
			if (comm.rank() == 0)      offset = {0,0};
			else if (comm.rank() == 1) offset = {5,0};
			else if (comm.rank() == 2) offset = {0,5};
			else if (comm.rank() == 3) offset = {5,5};
			return offset; 
		};

		SECTION( "Default initialisation, then assigning members" ){
			StructMeshPart<RectMesh> part;
			part.comm( comm );

			SECTION( "Assigning global mesh on master, then distributing" ) {
				auto assignGlobal = [&]( int ng ){

					RectMesh mesh;

					if( comm.isMaster() ){
						mesh = { createTrueRectMeshForMpi(ng) };
						addGlobalPatches(mesh);
					}

					part.assignAndDistribute( std::move(mesh) );

					test(part);
				};

				locallyAssigned = false;
				assignGlobal(1);
				part.clearBoundaryPatchLimits();
				assignGlobal(2);
			}

			SECTION( "Assigning local mesh" ){

				auto assignLocal = [&]( int ng ){
					RectMesh globalMesh { createTrueRectMeshForMpi(ng) };
					RectMesh localMesh {distributeMesh(globalMesh, comm, Array2i{10,10}
					,Array2i{5,5}, localOffset(), ng)};
					
					addLocalPatches(localMesh);

					part.assignLocal( std::move(localMesh), localOffset() );

					test(part);
				};
				locallyAssigned = true;
				assignLocal(1);
				part.clearBoundaryPatchLimits();
				assignLocal(2);
			}	
		}

		SECTION( "Assigning mesh during construction" ){
			SECTION( "Assigning global mesh on master" ){
				auto assignGlobal = [&](int ng){
					RectMesh mesh;

					if( comm.isMaster() ){
						mesh = createTrueRectMeshForMpi(ng);
						addGlobalPatches(mesh);
					}
					StructMeshPart part (std::move(mesh), comm);

					test(part);
				};
				locallyAssigned = false;
				assignGlobal(1);
				assignGlobal(2);
			}
			SECTION( "Assigning local meshes" ){
				auto assignLocal = [&]( int ng ){
					RectMesh globalMesh {createTrueRectMeshForMpi(ng) };
					RectMesh localMesh;

					localMesh = distributeMesh( globalMesh, comm, Array2i{10,10}, Array2i{5,5}, localOffset(), ng );

					localMesh.clearBoundaryPatches();
					addLocalPatches(localMesh);

					StructMeshPart part { std::move(localMesh), localOffset(), comm };

					test(part);
				};
				locallyAssigned = true;
				assignLocal(1);
				assignLocal(2);
			}
		}
	}
}

TEST_CASE( "Non-uniform RectMesh: Redistributing mesh" ){
	MPI_Comm_H comm{ MPI_COMM_WORLD};

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		auto test = [&]( int ng ){
			if (comm.size() != 4) return;

			auto [meshSizes0, meshSizes1, meshSizes2, offsets0, offsets1, offsets2]
				= sizesAndOffsetsForTrueRectMesh();

			RectMesh mesh;
			if( comm.isMaster() ){
				mesh = createTrueRectMeshForMpi(ng);
				addStandardBoundaryPatches(mesh);
			}

			StructMeshPart part { std::move(mesh), comm };

			REQUIRE( (part.meshSize() == meshSizes0.col( comm.rank() ) ).all() );
			REQUIRE( (part.offset()   == offsets0.col( comm.rank() ) ).all() );

			auto testStage1 = [&, &part = std::as_const(part)](){

				REQUIRE( (part.meshSize() == meshSizes1.col( comm.rank() ) ).all());
				REQUIRE( (part.offset()   == offsets1.col(   comm.rank() ) ).all() );
				REQUIRE( (part.oldSizes() == meshSizes0).all() );

				switch( comm.rank() ){
					case 0:
						REQUIRE( part.neighbours(BoundaryLocation::top)[0] == OverlapSegment{2,4,0,0} );
						REQUIRE( part.neighbours(BoundaryLocation::top)[1] == OverlapSegment{3,2,0,4} );
						REQUIRE( (part.mesh().boundaryPatch("bottom").edges() == ArrayXi::LinSpaced(6, 0, 6)).all() );
						REQUIRE( (part.mesh().boundaryPatch("left").edges()   == ArrayXi::LinSpaced(6, 42, 77)).all() );
						REQUIRE( part.mesh().boundaryPatch("right").empty());
						REQUIRE( part.mesh().boundaryPatch("top").empty());
						break;
					case 1:
						REQUIRE( part.neighbours(BoundaryLocation::top)[0] == OverlapSegment{3,4,2,0} );
						
						REQUIRE( (part.mesh().boundaryPatch("bottom").edges() == ArrayXi::LinSpaced(4, 0, 3)).all() );
						REQUIRE( (part.mesh().boundaryPatch("right").edges() == ArrayXi::LinSpaced(6, 32, 57)).all() );
						REQUIRE( part.mesh().boundaryPatch("left").empty());
						REQUIRE( part.mesh().boundaryPatch("top").empty());
						break;
					case 2:
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{0,4,0,0} );
											
						REQUIRE( part.mesh().boundaryPatch("right").empty() );
						REQUIRE( part.mesh().boundaryPatch("bottom").empty() );
						REQUIRE( (part.mesh().boundaryPatch("top").edges() == ArrayXi::LinSpaced(4, 16, 19)).all() );
						REQUIRE( (part.mesh().boundaryPatch("left").edges() == ArrayXi::LinSpaced(4, 20, 35)).all() );
						break;
					case 3: 
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{0,2,4,0} );
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] == OverlapSegment{1,4,0,2} );
											
						REQUIRE(  part.mesh().boundaryPatch("bottom").empty() );
						REQUIRE(  part.mesh().boundaryPatch("left").empty() );
						REQUIRE( (part.mesh().boundaryPatch("right").edges() == ArrayXi::LinSpaced(4, 36,57)).all() );
						REQUIRE( (part.mesh().boundaryPatch("top").edges() == ArrayXi::LinSpaced(6, 24, 29)).all() );
						break;
				}
			};

			auto testStage2 = [&, &part = std::as_const(part)](){
				REQUIRE( (part.meshSize() == meshSizes2.col( comm.rank() ) ).all() );
				REQUIRE( (part.offset() == offsets2.col( comm.rank() ) ).all() );
				REQUIRE( (part.oldSizes() == meshSizes1).all() );

				switch( comm.rank() ){
					case 0:
						REQUIRE( part.neighbours(BoundaryLocation::top)[0] ==	OverlapSegment{2, 3, 0, 0} );

						REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
						REQUIRE(part.mesh().boundaryPatch("right").empty());
						REQUIRE( (part.mesh().boundaryPatch("bottom").edges() == ArrayXi::LinSpaced(3, 0, 2) ).all() );
						REQUIRE( (part.mesh().boundaryPatch("left").edges() == ArrayXi::LinSpaced(4, 15, 27) ).all() );
						break;
					case 1:
						REQUIRE( part.neighbours(BoundaryLocation::top)[0] == OverlapSegment{2, 4, 3, 0} );
						REQUIRE( part.neighbours(BoundaryLocation::top)[1] == OverlapSegment{3, 3, 0, 4} );
						
						REQUIRE(part.mesh().boundaryPatch("top"  ).empty());
						REQUIRE(part.mesh().boundaryPatch("left").empty());
						REQUIRE( (part.mesh().boundaryPatch("bottom").edges() == ArrayXi::LinSpaced(7, 0, 6) ).all() );
						REQUIRE( (part.mesh().boundaryPatch("right").edges() == ArrayXi::LinSpaced(4, 42, 66) ).all() );
						break;
					case 2:
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{0, 3, 0, 0} );
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[1] == OverlapSegment{1, 4, 0, 3} );
					
						REQUIRE(part.mesh().boundaryPatch("bottom").empty());
						REQUIRE(part.mesh().boundaryPatch("right").empty());
						REQUIRE( (part.mesh().boundaryPatch("top").edges() == ArrayXi::LinSpaced(7, 42, 48) ).all() );
						REQUIRE( (part.mesh().boundaryPatch("left").edges() == ArrayXi::LinSpaced(6, 49, 89) ).all() );
						break;
					case 3:
						REQUIRE( part.neighbours(BoundaryLocation::bottom)[0] == OverlapSegment{1, 3, 4, 0} );
						
						REQUIRE(part.mesh().boundaryPatch("bottom").empty());
						REQUIRE(part.mesh().boundaryPatch("left").empty());
						REQUIRE( (part.mesh().boundaryPatch("top").edges() == ArrayXi::LinSpaced(3, 18, 20) ).all() );
						REQUIRE( (part.mesh().boundaryPatch("right").edges() == ArrayXi::LinSpaced(6, 24, 44) ).all() );
						break;
				}
			};

			SECTION( "Pure Redistribution - " + std::to_string(ng) + " ghost layers" ){
				part.redistributeMesh( meshSizes1 );

				SECTION( "Single redistribution" ){
					testStage1();
				}
				SECTION( "Repeated Distribution" ){
					part.redistributeMesh( meshSizes2 );
					testStage2();
				}
			}
		};
		test(1);
		test(2);
	}
}


TEST_CASE( "Non-uniform RectMesh: (Re-)Distributing field" ){
	MPI_Comm_H comm {MPI_COMM_WORLD};

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		if( comm.size() == 4){
			auto [meshSizes0, meshSizes1, meshSizes2, offsets0, offsets1, offsets2] = sizesAndOffsetsForTrueRectMesh();

			RectMesh mesh;
			Field<RectMesh> f;

			auto writeRanksToField = [&](int nComponents, int ng, const Ref<const Array2Xi>& offsets, const Ref<const Array2Xi>& meshSizes){
				
				if( comm.isMaster() ){
					mesh = createTrueRectMeshForMpi(ng);
					f = Field( mesh, nComponents );
					f.values() = 0.99;
					for( int i{0}; i < comm.size(); i++) 
						f.values().reshaped(
							(mesh.nCellsX() + 2*ng)*nComponents, 
							mesh.nCellsY() + 2*ng
							).block(
								(offsets(0,i)+ng)*nComponents, offsets(1,i)+ng, 
								meshSizes(0,i)*nComponents, meshSizes(1,i)
								) = i;
				}
			};

			SECTION( "Distributing field" ) {
				auto test = [&]( int nComponents, int ng ){
					writeRanksToField( nComponents, ng, offsets0, meshSizes0 );
					
					StructMeshPart part ( std::move(mesh), comm );

					part.distributeField( f );

					REQUIRE( (f.domainValues() == comm.rank()).all() );
				};

				for ( int nVals {1}; nVals<4; ++nVals )
					for ( int ng {1}; ng<3; ++ng )
						test( nVals, ng );
			}
			SECTION( "Redistributing field once" ){
				auto test = [&]( int nComponents, int ng ){
					writeRanksToField( nComponents, ng, offsets1, meshSizes1 );

					StructMeshPart part ( std::move(mesh), comm );
					part.distributeField(f);
					part.redistributeMesh( meshSizes1 );
					part.redistributeField(f);
					
					REQUIRE( ( f.domainValues() == comm.rank() ).all() );
				};
				for ( int nVals {1}; nVals<4; ++nVals )
					for ( int ng {1}; ng<3; ++ng )
						test( nVals, ng );
			}
			SECTION( "Repeatedly redistributing field" ){
				auto test = [&](int nComponents, int ng){
					writeRanksToField( nComponents, ng, offsets2, meshSizes2 );

					StructMeshPart part ( std::move(mesh), comm );
					part.distributeField(f);
					part.redistributeMesh( meshSizes1 );
					part.redistributeField(f);
					part.redistributeMesh( meshSizes2 );
					part.redistributeField(f);
					
					REQUIRE( ( f.domainValues() == comm.rank() ).all() );
				};
				for ( int nVals {1}; nVals<4; ++nVals )
					for ( int ng {1}; ng<3; ++ng )
						test( nVals, ng );
			}
		}
	}
}

TEST_CASE( "Non-uniform RectMesh: getNeighbourValues: Retrieving domain values from neighbours to update ghost values" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		auto test = [&]( Index ng, Index nComponents ){

			RectMesh mesh;

			StructMeshPart part { comm.isMaster() ? createTrueRectMeshForMpi(ng) : std::move(mesh), comm };

			Field f { part.mesh(), nComponents };

			auto testField = [&](){

				f.values() = -1;

				f.domainValues() = comm.rank();
				part.getNeighbourValues( f );

				auto valsReshaped { f.values().reshaped(
					(part.meshSize().x() + 2 * ng)*nComponents,
					part.meshSize().y() + 2 * ng
				) };

				auto bottomVals { valsReshaped.block(
					ng*nComponents, 0,
					part.meshSize().x()*nComponents, ng
				) };
				auto topVals { valsReshaped.block(
					ng*nComponents, ng + part.meshSize().y(),
					part.meshSize().x()*nComponents, ng
				) };
				auto leftVals { valsReshaped.block(
					0, ng, ng*nComponents, part.meshSize().y()
				) };
				auto rightVals { valsReshaped.block(
					(ng + part.meshSize().x())*nComponents, ng,
					ng*nComponents, part.meshSize().y()
				) };

				for ( const auto& nb : std::as_const(part).neighbours(BoundaryLocation::bottom) )
					REQUIRE( bottomVals.middleRows(
						nb.ownBeg*nComponents, nb.size*nComponents
					).isApproxToConstant(nb.rank, epsilon) );
				
				for ( const auto& nb : std::as_const(part).neighbours(BoundaryLocation::top) )
					REQUIRE( ( topVals.middleRows(
						nb.ownBeg*nComponents, nb.size*nComponents
					) == nb.rank ).all() );
				
				if ( !part.atBoundary(BoundaryLocation::left) )
					REQUIRE( (leftVals == comm.rank()-1).all() );
				
				if ( !part.atBoundary(BoundaryLocation::right) )
					REQUIRE( (rightVals == comm.rank()+1).all() );
			};

			testField();

			ArrayXs load ( comm.size() );
			for ( int i{0}; i<comm.size(); ++i ){
				load(i) = comm.size() + i;
			}
			part.rebalanceMeshSizes( PartLoad::perNode, load );

			part.redistributeField(f);

			testField();
		};
		for ( int nVals {1}; nVals<4; ++nVals )
			for ( int ng {1}; ng<3; ++ng )
				test( nVals, ng );
	}
}

TEST_CASE( "Non-uniform RectMesh: nRanks invariance, construction" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	if ( comm.size() != 4 ){
		if ( comm.isMaster() )
			std::cerr << "WARNING: tests relating to true StructMeshes and RectMeshes are only "
			"executed if the number of processes is set to 4. Otherwise they will be skipped.";
	}
	else{
		int ng {2};
		Array2i globalMeshSize{10,10};
		Array2Xi patchLimits (2, 4);
		patchLimits <<
			0, 110, 120, 100,
			9, 209, 219, 109;

		StructMeshPart<RectMesh> part;

		// auto print = [&]( const auto& message ){
		// 	std::stringstream sstream;
		// 	sstream << "Rank " << comm.rank() << ": " << message << "\n";
		// 	for (int i{0};i<comm.size();++i){
		// 		MPI_Barrier(comm);
		// 		if (comm.rank()==i)
		// 			std::cerr << sstream.str();
		// 	}
		// };	

		auto test = [&](){
			REQUIRE( comm.size() == part.comm().size() );
			REQUIRE( comm.rank() == part.comm().rank() );
			REQUIRE( (part.meshSizeGlobal() == globalMeshSize).all() );
			REQUIRE( (part.boundaryPatchLimits() == patchLimits).all() );
			REQUIRE( part.meshSize().x() == globalMeshSize.x() / part.nParts().x() +
				(part.indexInBand() == part.nParts().x() - 1 ?
				globalMeshSize.x() % part.nParts().x() : 0)
			);
			REQUIRE( part.meshSize().y() == globalMeshSize.y() / part.nParts().y() +
				(part.band() == part.nParts().y() - 1 ?
				globalMeshSize.y() % part.nParts().y() : 0)
			);
		};

		part.comm( comm );

		RectMesh mesh;

		if( comm.isMaster() ){
			mesh = createTrueRectMeshForMpi(ng);
			addStandardBoundaryPatches(mesh);
		}

		part.assignAndDistribute( std::move(mesh) );

		test();

		Array2i localSize {5,5};
		Array2i offset;
		if(comm.rank() == 0)      offset = {0,0};	
		else if(comm.rank() == 1) offset = {5,0};
		else if(comm.rank() == 2) offset = {0,5};	
		else if(comm.rank() == 3) offset = {5,5};

		mesh = distributeMesh(mesh, comm, globalMeshSize, localSize, offset, ng);
		mesh.clearBoundaryPatches();

		mesh.addBoundaryPatches( part.mesh().moveBoundaryPatches() );
		part = { std::move(mesh), part.offset(), comm };

		test();

		Array2i newSize { part.meshSize() - 1 };
		if ( part.indexInBand() == part.nParts().x() - 1 )
			newSize.x() += part.nParts().x();
		if ( part.band() == part.nParts().y() - 1 )
			newSize.y() += part.nParts().y();

		Array2Xi redistSizes ( 2, comm.size() );
		MPI_Allgather(
			newSize.data(), 2, MPI_INDEX,
			redistSizes.data(), 2, MPI_INDEX, comm
		);

		part.redistributeMesh(redistSizes);
	}
}